import torch
import whisper
import time
#首先清空运算缓存
torch.cuda.empty_cache()
#加载模型使用load_model方法
'''
load_model有三个参数：
    第一个参数是需要加载的模型名称，模型名称分为tiny、base、small、medium、large。
    第二个参数是计算使用的设备，设备可以使用cup也可以使用gpu,如果安装了CUDA则默认使用GPU,否则会默认使用CPU。
        使用CPU: torch.device("cpu")    使用GPU: torch.device("cuda")
    第三个参数是模型所在路径，如果为空，则会放在~/.cache/whisper下。
'''
whisper_model = whisper.load_model("medium",torch.device("cuda"),r"D:\AI\Whisper")
audioFile = input("请输入文件路径：")
startTime = time.time() #记录开始时间
#加载模型
'''
transcribe可以接受很多参数
    第一个参数是要翻译的录音地址，最重要的就是第一个参数，其他参数都可选。
    可以使用language指定对应的语言。
    其他的参数能够调整语速等设置，能够使识别更准确，但我还没有深入研究
    initial_prompt='阳光下的泡沫，是彩色的' 初始的提示，对于whisper理解音频有帮助。
    response_format，也就是返回的文件格式，我们这里是默认值，也就是 JSON。实际你还可以选择 TEXT 这样的纯文本，或者 SRT 和 VTT 这样的音频字幕格式。这两个格式里面，除了文本内容，还会有对应的时间信息，方便你给视频和音频做字幕。你可以直接试着运行一下看看效果。
    temperature，这个和我们之前在 ChatGPT 类型模型里的参数含义类似，就是采样下一帧的时候，如何调整概率分布。这里的参数范围是 0-1 之间。
'''
result = whisper_model.transcribe(audioFile,language='ja')
#全部输出，使用逗号分割
print(", ".join([i["text"] for i in result["segments"] if i is not None]))
#标注时间并逐句输出
# for word in result['segments']:
#     print('【',word['start'],'-',word['end'],'】:',word['text'])
endTime = time.time() #记录结束时间
print("程序运行时间为：",endTime-startTime)